package agent;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.Collections;
import com.alibaba.dashscope.app.Application;
import com.alibaba.dashscope.app.ApplicationParam;
import com.alibaba.dashscope.app.ApplicationResult;
import com.alibaba.dashscope.exception.ApiException;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.alibaba.fastjson.JSON;
import common.MessageMissionInfo;
import org.apache.rocketmq.client.apis.ClientConfiguration;
import org.apache.rocketmq.client.apis.ClientException;
import org.apache.rocketmq.client.apis.ClientServiceProvider;
import org.apache.rocketmq.client.apis.SessionCredentialsProvider;
import org.apache.rocketmq.client.apis.StaticSessionCredentialsProvider;
import org.apache.rocketmq.client.apis.consumer.ConsumeResult;
import org.apache.rocketmq.client.apis.consumer.FilterExpression;
import org.apache.rocketmq.client.apis.consumer.FilterExpressionType;
import org.apache.rocketmq.client.apis.message.Message;
import org.apache.rocketmq.client.apis.producer.Producer;
import org.apache.rocketmq.client.apis.producer.ProducerBuilder;
import org.apache.rocketmq.client.apis.producer.SendReceipt;
import org.apache.rocketmq.client.apis.producer.TransactionChecker;
import org.apache.rocketmq.shaded.commons.lang3.StringUtils;

public class WeatherAgentMain {
    private static final String GlobalEndpoint = System.getProperty("endpoint");
    private static final String RocketMQInstanceID = System.getProperty("rocketmqInstanceID");
    private static final String WorkAgentResponseTopic = System.getProperty("workAgentResponseTopic");
    private static final String BizTopic = System.getProperty("weatherAgentTaskTopic");
    private static final String BizConsumerGroup = System.getProperty("weatherAgentTaskGroupID");
    private static final String AccessKey = System.getProperty("ak");
    private static final String SecretKey = System.getProperty("sk");
    private static final String ApiKey = System.getProperty("apiKey");
    private static final String AppId = System.getProperty("appId");
    private static final String AgentName = "WeatherAgent";

    public static void main(String[] args) throws ClientException, IOException, InterruptedException {
        if (!checkConfigParam()) {
            System.out.println("配置参数不完整，请检查参数配置情况");
            return;
        }
        init();
        System.out.println("启动" + AgentName + "成功，开始监听消息, Topic: " + BizTopic);
    }

    private static boolean checkConfigParam() {
        if (StringUtils.isEmpty(GlobalEndpoint) || StringUtils.isEmpty(RocketMQInstanceID)
            || StringUtils.isEmpty(WorkAgentResponseTopic) || StringUtils.isEmpty(BizTopic) || StringUtils.isEmpty(BizConsumerGroup)
            || StringUtils.isEmpty(AccessKey) || StringUtils.isEmpty(SecretKey)
            || StringUtils.isEmpty(ApiKey) || StringUtils.isEmpty(AppId)) {
            if (StringUtils.isEmpty(GlobalEndpoint)) {
                System.out.println("请配置RocketMQV5的实例接入点 endpoint");
            }
            if (StringUtils.isEmpty(RocketMQInstanceID)) {
                System.out.println("请配置RocketMQV5的实例信息 rocketmqInstanceID");
            }
            if (StringUtils.isEmpty(WorkAgentResponseTopic)) {
                System.out.println("请配置RocketMQV5的轻量消息Topic workAgentResponseTopic");
            }
            if (StringUtils.isEmpty(BizTopic)) {
                System.out.println("请配置RocketMQV5的 weatherAgentTaskTopic");
            }
            if (StringUtils.isEmpty(BizConsumerGroup)) {
                System.out.println("请配置RocketMQV5的 weatherAgentTaskGroupID");
            }
            if (StringUtils.isEmpty(AccessKey)) {
                System.out.println("请配置RocketMQV5的访问控制-用户名 ak");
            }
            if (StringUtils.isEmpty(SecretKey)) {
                System.out.println("请配置RocketMQV5的访问控制-密码 sk");
            }
            if (StringUtils.isEmpty(ApiKey)) {
                System.out.println("请配置天气助手Agent在百炼平台的 apiKey");
            }
            if (StringUtils.isEmpty(AppId)) {
                System.out.println("请配置天气助手Agent在百炼平台的 appId");
            }
            return false;
        }
        return true;
    }

    private static void init() throws ClientException {
        final ClientServiceProvider provider = ClientServiceProvider.loadService();
        SessionCredentialsProvider sessionCredentialsProvider = new StaticSessionCredentialsProvider(AccessKey, SecretKey);
        ClientConfiguration clientConfiguration = ClientConfiguration.newBuilder()
            .setEndpoints(GlobalEndpoint)
            .setNamespace(RocketMQInstanceID)
            .setCredentialProvider(sessionCredentialsProvider)
            .build();
        String tag = "*";
        FilterExpression filterExpression = new FilterExpression(tag, FilterExpressionType.TAG);
        Producer producer = buildProducer(null, WorkAgentResponseTopic);
        provider.newPushConsumerBuilder()
            .setClientConfiguration(clientConfiguration)
            .setConsumerGroup(BizConsumerGroup)
            .setSubscriptionExpressions(Collections.singletonMap(BizTopic, filterExpression))
            .setMessageListener(messageView -> {
                try {
                    printSystemSuccess(AgentName + " 开始拉取到普通消息, Topic: " + BizTopic + ", messageId: " + messageView.getMessageId());
                    byte[] result = new byte[messageView.getBody().remaining()];
                    messageView.getBody().get(result);
                    String messageStr = new String(result, StandardCharsets.UTF_8);
                    MessageMissionInfo messageMissionInfo = JSON.parseObject(messageStr, MessageMissionInfo.class);
                    String response = appCall(messageMissionInfo.getMessageText());
                    System.out.println("大模型响应结果: " + response);
                    Message message = buildMessage(WorkAgentResponseTopic, messageMissionInfo.getLiteTopic(), response, messageMissionInfo.getMessageText());
                    final SendReceipt sendReceipt = producer.send(message);
                    printSystemSuccess("将大模型响应结果发送到 " + messageMissionInfo.getRequestAgentName() + " 订阅的 Topic: " + WorkAgentResponseTopic + ", LiteTopic: " + messageMissionInfo.getLiteTopic() + " 成功, messageId: " + sendReceipt.getMessageId());
                } catch (Exception e) {
                    System.out.println("error" + e.getMessage());
                    return ConsumeResult.FAILURE;
                }
                return ConsumeResult.SUCCESS;
            }).build();
    }

    private static Message buildMessage(String topic, String liteTopic, String response, String request) {
        if (StringUtils.isEmpty(topic) || StringUtils.isEmpty(liteTopic) || StringUtils.isEmpty(response)  || StringUtils.isEmpty(request)) {
            System.out.println("buildMessage param error topic: " + topic + ", liteTopic: " + liteTopic + ", response: " + response + "request: " + request);
            return null;
        }
        MessageMissionInfo messageMissionInfo = new MessageMissionInfo();
        messageMissionInfo.setMessageText(response);
        messageMissionInfo.setMessageRequest(request);
        messageMissionInfo.setResponseAgentName(AgentName);
        String missionJsonStr = JSON.toJSONString(messageMissionInfo);
        final ClientServiceProvider provider = ClientServiceProvider.loadService();
        final Message message = provider.newMessageBuilder()
            .setTopic(topic)
            .setBody(missionJsonStr.getBytes(StandardCharsets.UTF_8))
            .setLiteTopic(liteTopic)
            .build();
        return message;
    }

    private static Producer buildProducer(TransactionChecker checker, String... topics) throws ClientException {
        final ClientServiceProvider provider = ClientServiceProvider.loadService();
        SessionCredentialsProvider sessionCredentialsProvider = new StaticSessionCredentialsProvider(AccessKey, SecretKey);
        ClientConfiguration clientConfiguration = ClientConfiguration.newBuilder()
            .setEndpoints(GlobalEndpoint)
            .setNamespace(RocketMQInstanceID)
            .setCredentialProvider(sessionCredentialsProvider)
            .setRequestTimeout(Duration.ofSeconds(10))
            .build();
        final ProducerBuilder builder = provider.newProducerBuilder()
            .setClientConfiguration(clientConfiguration)
            .setTopics(topics);
        if (checker != null) {
            builder.setTransactionChecker(checker);
        }
        return builder.build();
    }

    public static String appCall(String prompt) throws ApiException, NoApiKeyException, InputRequiredException {
        ApplicationParam param = ApplicationParam.builder()
            .apiKey(ApiKey)
            .appId(AppId)
            .prompt(prompt)
            .build();
        ApplicationResult result = new Application().call(param);
        return result.getOutput().getText();
    }

    private static void printSystemSuccess(String message) {
        System.out.println("\u001B[32m[SUCCESS] " + message + "\u001B[0m");
    }

    private static void printSystemSuccessSimple(String message) {
        System.out.print("\u001B[32m[SUCCESS] " + message + "\u001B[0m");
    }

}
